#############################################################################
# 本脚本展示使用PyTorch从零开始实现Transformer模型的Encoder部分, 代码参考如下连接:
# https://www.bilibili.com/video/BV13h411c7hW/?p=25&spm_id_from=pageDriver
#
# 
#                                                       Author： Xian Yang
#                                                     Date Time： 2023.11.25
##############################################################################

import torch
import copy
import math                                   # 导入数学工具包
import numpy as np
import torch.nn as nn                         # torch.nn预定义了一些网络层，我们可以直接使用
import torch.nn.functional as F
import torch.autograd.variable as Variable    # torch中的变量封装函数

from   typing import Tuple, Map, Union

class Embedding(nn.Module):
    """_summary_

    Args:
        vacab   (int): 词汇表的大小
        d_model (int): 词嵌入的大小，或者说隐藏层的大小
    """
    def __init__(self, vocab : int, d_model : int, padding_idx: int = 0) -> None:
        """
        当你使用 padding_idx 参数时，这里是一些关键点：
        1. 填充的目的：在处理文本或序列数据时，不同的样本通常有不同的长度。为了能够以固定大小的批处理这些数据，
            较短的序列通常会用某个特定的“填充”值扩充到与最长序列相同的长度。

        2. padding_idx 参数：在 nn.Embedding 层中设置 padding_idx 参数后，该索引对应的嵌入向量将被初始
            化为零，并且在训练过程中不会被更新。这意味着无论该填充索引出现多少次，其嵌入向量始终为零向量。

        3.对模型的影响：使用 padding_idx 可以确保填充的索引不会影响到模型的学习过程。由于其嵌入向量始终为零，
            因此填充的部分在进行任何形式的向量运算时，都不会对结果产生实质性的影响。

        例如，如果你设定 padding_idx=0，那么索引为0的词汇项（通常用于填充较短的序列）的嵌入向量在模型初始化时
            被设置为零向量，并且在后续的训练过程中保持不变。
        """
        super(Embedding, self).__init__()      # 显式调用父类的构造函数
        self._lut     = nn.Embedding(vocab, d_model, padding_idx)
        self._d_model = d_model

    def forward(self, x):                      # 起到一个缩放的作用
        return self._lut(x) * math.sqrt(self._d_model)
    
    @property
    def hidden_size(self):
        return self._d_model
    

class PositionEmbedding(nn.Module):
    def __init__(self, d_model : int, dropout : float, max_len : int = 5000) -> None:
        r"""_summary_

        Args:
            d_model (int):           词嵌入(隐藏层)大小
            dropout (float):         置零的比例
            max_len (int, optional): 句子最大长度, 默认为5000
        """
        super(PositionEmbedding, self).__init__()
        self._dropout = nn.Dropout(p = dropout)

        # 初始化一个位置编码矩阵
        pe            = torch.zeros(max_len, d_model, dtype=torch.float32)

        pos           = torch.arange(0, max_len, dtype=torch.float32).reshape(-1, 1) 
        div_term      = torch.exp(torch.arange(0, d_model, 2) * -(math.log(10000.0)/d_model))
        pe[:, 0::2]   = torch.sin(pos * div_term)
        pe[:, 1::2]   = torch.cos(pos * div_term)
        pe            = pe.unsqueeze(0)           # 在第0维增加一个维度，为了与输入维度匹配

        # TODO: 特别注意
        # 把pe位置编码矩阵注册成模型的buffer，什么是模型的buffer呢？
        # 我们认为它是对模型有帮助的，但是不是模型中的参数或超参, 其不需要随模型的训练而更新
        # 注册之后我们可以在模型保存后重新加载时和模型结构与参数一同被加载
        self.register_buffer('pe', pe)

    def forward(self, x : torch.Tensor) -> torch.Tensor:  
        #  x.shape = [batch, seq_len, d_model]
        x = x + Variable(self.pe[:, x.size(1)], required=False)

        # 使用dropout对象进行丢弃操作，返回结果
        # dropput的作用是使用指定的概率将输入的tensor置零
        r"""pyhton
            >>> import torch
            >>> import torch.nn as nn
            >>> m = nn.Dropout(p=0.2)
            >>> x = torch.randn(4, 5)
            >>> y = m(x)
            >>> y
            tensor([[ 0.0000, -1.3313,  0.0000,  0.0000, -0.0000],
                    [-0.0000, -2.0547, -0.4483,  1.7132, -0.6219],
                    [-1.2918, -0.0026, -0.0000, -1.1671,  0.2106],
                    [ 0.1137,  0.2830,  0.0770, -2.2582,  0.0000]])
        """
        return self._dropout(x)
    

class MultiHeadAttention(nn.Module):
    def __init__(self, num_head : int, d_model : int, dropout: float = 0.1) -> None:
        """_summary_
            多有自注意力机制的类
        Args:
            num_head (int):   多头的数量，
            d_model  (int):   词嵌入维度
            dropout  (float): dropout置0的概率
        """
        super(MultiHeadAttention, self).__init__()
        assert d_model % num_head == 0, "d_model a should be divisible by num_head"
        self.d_k      = d_model // num_head
        self.num_head = num_head

        # TODO: 为什么？
        # 多头注意力机制中，Q,K,V各有一个权重矩阵，Concate后还有一个线性矩阵
        # 理论上应该是 num_head * 4个(d_model, d_model / num_head)的矩阵，这里相当于拼在一起了
        # TODO: 这里就不能按普通的矩阵乘法，而需要使用batch的矩阵乘法
        self.linears  = MultiHeadAttention.clones(nn.Linear(d_model, d_model), 4)
        self.attn     = None
        self.dropout  = nn.Dropout(dropout) 


    # 在多头自注意力机制中，要多个多个结构相同的线性层
    # 我们定一个克隆函数，将他们一同初始化在网络层列表对象中
    @staticmethod
    def clones(module : nn.Module, num : int) -> nn.ModuleList:
        return nn.ModuleList([copy.deepcopy(module) for _ in range(num)])
    
    # 三维张量: attn_shape = (1, max_len, max_len)
    @staticmethod
    def gen_mask(max_len : int) -> torch.Tensor:
        """_summary_
            生成向后这样的掩码张量, 参数size是
        Args:
            size (int): 掩码张量的尺寸
        """
        attn_shape = (1, max_len, max_len)
        """_summary_
            >>> size = 5
            >>> attn_shape = (1, max_len, max_len)
            # k = 0 表示主对角线, k > 0 表示向右上移动
            >>> mask       = np.triu(np.ones(shape=attn_shape), k = 1).astype(np.uint8)
            >>> mask
                array([[[0, 1, 1, 1, 1],
                        [0, 0, 1, 1, 1],
                        [0, 0, 0, 1, 1],
                        [0, 0, 0, 0, 1],
                        [0, 0, 0, 0, 0]]], dtype=uint8)
            >>> torch.from_numpy(1 - mask)
                tensor([[[1, 0, 0, 0, 0],
                        [1, 1, 0, 0, 0],
                        [1, 1, 1, 0, 0],
                        [1, 1, 1, 1, 0],
                        [1, 1, 1, 1, 1]]], dtype=torch.uint8)
        """
        # TODO: 向后遮掩
        mask   = np.triu(np.ones(shape=attn_shape), k = 1).astype(np.uint8)
        return torch.from_numpy(1 - mask)


    @staticmethod
    def attention(query : torch.Tensor, key : torch.Tensor, value : torch.Tensor, 
                    mask : torch.Tensor = None, dropout :  nn.Dropout = None
                ) -> Tuple[torch.Tensor, torch.Tensor]:
        """_summary_

        Args:
            query   (torch.Tensor):  查询矩阵      (batch_size, num_head, max_len, d_k)
            key     (torch.Tensor):  key矩阵      (batch_size, num_head, max_len, d_k)
            value   (torch.Tensor):  value矩阵    (batch_size, num_head, max_len, d_k)
            mask    (torch.Tensor):  掩码矩阵
            dropout (int, optional): 置零

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: _description_
        """
        d_k      = query.size(-1)        #  hidden_size / head_num, 即size per head
        # 按照公式将query，与key的转置相乘(这里转置key的最后两个维度)，然后在归一化
        scores   = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)

        # 判断是否使用掩码张量
        # 当 mask 中的元素等于0时，它会将 scores 张量中对应位置的值替换为-1e9。
        # TODO: 为什么用非常小的数值？ 小数值表示word之间的attn很小
        if mask is not None:
            scores = scores.masked_fill(mask == 0, -1e9)

        # 对最后一个维度进行softmax, 这里其实就是按行
        p_attn = F.softmax(scores, dim = -1)    
        if dropout is not None:
            p_attn = dropout(p_attn)

        # 最终返回注意力矩阵，同时返回注意力张量
        return torch.matmul(p_attn, value), p_attn


    def forward(self, query : torch.Tensor, key : torch.Tensor, value : torch.Tensor, 
                mask : torch.Tensor = None) -> torch.Tensor:

        if mask is not None:
            mask = mask.unsqueeze(1)               # 扩展维度，代码多头中的第n头
        
        batch_size = query.size(0)                 # batch的大小，也即样本数量

        # 使用线性变换分别计算出QKV的值
        # 这里的-1其实PyTorch自适应维度,最后应该得到的值为max_len
        # transpose的目的是为了能够让代表句子长度的的维度能够与词向量的维度相邻
        query, key, value = [model(x).view(batch_size, -1, self.num_head, self.d_k).transpose(1, 2) 
                             for model, x in zip(self.linears, (query, key, value))]
        
        # 计算多头注意力输出
        x, self.attn = MultiHeadAttention.attention(query, key, value, mask, dropout=self.dropout)

        # 接下来进行concate, 将维度变成: (batch_size, max_len, d_model)
        # contiguous() 方法的本质是让转置后的内存连续而不是一个视图
        # 后面再进行view的时候，必须要先调用contiguous()，这是PyTorch的一个限制
        x = x.transpose(1,2).contiguous().view(batch_size, -1, self.num_head, self.d_k)
 
        # 最后一个线性层将concate后的数据进行线性变换
        return self.linears[-1](x)     


class FeedForward(nn.Module):
    def __init__(self, d_model : int, d_ff : int , dropout : float = 0.1) -> None:
        """_summary_
            前馈全连接层
        Args:
            d_model (int):    输入层维度
            d_ff (int):       隐藏层维度
            dropout (float):  置零概率
        """
        super(FeedForward, self).__init__()
        self.fft1    = nn.Linear(d_model, d_ff)
        self.fft2    = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)


    def forward(self, x):
        # 第一层有激活函数，第二层没有, 原算法就是这样的
        return self.fft2(self.dropout(F.relu(self.fft1(x))))


class LayerNorm(nn.Module):
    def __init__(self, d_model :int, eps : float = 1e-6) -> None:
        """_summary_
            LayerNorm的实现，注意LayerNorm是要归一化单个句子，
        Args:
            d_model (int):   
            eps     (float): 防止归一化时分母为0. Defaults to 1e-6.
        """
        super(LayerNorm, self).__init__()  
        # nn.Parameter 表示这个是可训练的参数
        self._alpha = nn.Parameter(torch.ones(d_model))
        self._beta  = nn.Parameter(torch.zeros(d_model))
        self._eps   = eps

    
    def forward(self, x):
        # x.shape:  (batch_size, max_len, d_model)
        _mean = x.mean(-1, keepdim=True)
        _std  = x.std(-1,  keepdim=True)
        return self._alpha * (x - _mean) / (_std + self._eps) + self._beta


class ShortCut(nn.Module):
    def __init__(self, d_model :int, eps : float = 1e-6, dropout : float = 0.1) -> None:
        super(ShortCut, self).__init__()
        self._layernorm = LayerNorm(d_model = d_model, eps = eps)
        self._dropout   = nn.Dropout(p=dropout)

    def forward(self, x : torch.Tensor, sublayer : Union[FeedForward, MultiHeadAttention]) -> torch.Tensor:
        # https://www.bilibili.com/video/BV13h411c7hW/?p=25&spm_id_from=pageDriver&vd_source=755665b659d6111b9b5a3247d5988a1d
        # 时间点: 2分21秒
        # return x + self._dropout(sublayer(self._layernorm(x)))
        # TODO: 是这样的么？难道不是下面这样？我认为我是正确的，所以先注释上述部分
        return self._layernorm(x + self._dropout(sublayer(x)))



class EncoderBlock(nn.Module):
    def __init__(self, d_model : int, attn : MultiHeadAttention,
                ffn : FeedForward, dropout : float = 0.1)->None:
        super(EncoderBlock, self).__init__()
        self._attn      = attn
        self._ffn       = ffn
        self._short_cut = MultiHeadAttention.clones(ShortCut(d_model, eps=1e-6, dropout=dropout), 2)
        self._d_model   = d_model

    def forward(self, x, mask):
        x = self._short_cut[0](x, lambda x: self._attn(x, x, x,mask))
        return self._short_cut[1](x, self._ffn)
        

class Encoder(nn.Module):
    def __init__(self, layer, num_block) -> None:
        super().__init__()
        self._layers = MultiHeadAttention.clones(layer, num_block)
        self._layernorm = LayerNorm(layer.size)

    def forward(self, x, mask):
        for layer in self._layers:
            x = layer(x, mask)
        return self._layernorm(x)